#!/usr/bin/env python
# encoding=utf-8

from collections import defaultdict
import json

COUNTRIES_FILE = 'countries.json'
INDEX_FILE = 'countriesIndex.json'


def createIndex(countriesFile=COUNTRIES_FILE, indexFile=INDEX_FILE):
    """ Create index for quick country search.
    """
    with open(countriesFile) as f:
        countries = json.load(f)

    index = defaultdict(list)
    for iso in countries:
        country = countries[iso]
        name = country['name'].lower().replace(' ', '')
        for i in range(len(name)-2):
            trigram = name[i:i+3]
            index[trigram].append(iso)

    with open(indexFile, 'w') as f:
        f.write(json.dumps(index))


if __name__ == "__main__":
    createIndex()
